/* routines that supply a kernel heap */

#include	"kheap.h"
#include	"monitor.h"
#include	"paging.h"

extern u32int end;	/* defined in link.ld - end addr of kernel text/bss/data */
extern struct page_directory *kernel_directory;
extern struct page_directory *current_directory;
u32int placement_addr = (u32int)&end;
struct heap *kheap = (struct heap *)0;
struct heap *uheap = (struct heap *)0;

u32int kmalloc_int(u32int sz, int align, u32int *phys){

	u32int tmp;
	struct page *page;
	void *addr;

	if(kheap){
		addr = alloc(sz, (u8int)align, kheap);

		if(phys){
			page = get_page((u32int)addr, 0, kernel_directory);
			*phys = (u32int)(page -> frame * 0x1000) + ((u32int)addr & 0xFFF);
		}
		return (u32int)addr;
	}else{
		if(align && (placement_addr & 0xFFFFF000)){
			/* not a page-aligned addr */
			placement_addr &= 0xFFFFF000;
			placement_addr += 0x1000;
		}
		if(phys){
			/* returns a virtual addr, save the physical addr in *phys */
			*phys = placement_addr;
		}
		/* stub */
		tmp = placement_addr;
		placement_addr += sz;

		return tmp;
	}
}

u32int kmalloc(u32int sz){
	return kmalloc_int(sz, 0, (u32int *)0);
}

u32int kmalloc_a(u32int sz){
	return kmalloc_int(sz, 1, (u32int *)0);
}

u32int kmalloc_p(u32int sz, u32int *phys){
	return kmalloc_int(sz, 0, phys);
}

u32int kmalloc_ap(u32int sz, u32int *phys){
	return kmalloc_int(sz, 1, phys);
}

void kfree(void *p){
	free(p, kheap);
}

u32int umalloc_int(u32int sz, int align, u32int *phys){

	u32int tmp;
	struct page *page;
	void *addr;

	if(uheap){
		addr = alloc(sz, (u8int)align, uheap);

		if(phys){
			page = get_page((u32int)addr, 0, current_directory);
			*phys = (u32int)(page -> frame * 0x1000) + ((u32int)addr & 0xFFF);
		}
		return (u32int)addr;
	}else{
		if(align && (placement_addr & 0xFFFFF000)){
			/* not a page-aligned addr */
			placement_addr &= 0xFFFFF000;
			placement_addr += 0x1000;
		}
		if(phys){
			/* returns a virtual addr, save the physical addr in *phys */
			*phys = placement_addr;
		}
		/* stub */
		tmp = placement_addr;
		placement_addr += sz;

		return tmp;
	}
}

u32int umalloc(u32int sz){
	return umalloc_int(sz, 0, (u32int *)0);
}

u32int umalloc_a(u32int sz){
	return umalloc_int(sz, 1, (u32int *)0);
}

u32int umalloc_p(u32int sz, u32int *phys){
	return umalloc_int(sz, 0, phys);
}

u32int umalloc_ap(u32int sz, u32int *phys){
	return umalloc_int(sz, 1, phys);
}

void ufree(void *p){
}

static void expand(u32int size, struct heap *pheap){

	u32int i, oldsize;

	if(size & 0xFFFFF000){
		size &= 0xFFFFF000;
		size += 0x1000;
	}
	oldsize = pheap -> end_address - pheap -> start_address;
	for(i = oldsize; i < size; i += 0x1000)
		alloc_frame(get_page(pheap -> start_address + i, 1, kernel_directory), (pheap -> supervisor) ? 1 : 0, (pheap -> readonly) ? 0 : 1);
	pheap -> end_address = pheap -> start_address + size;
}

static u32int contract(u32int size, struct heap *pheap){

	u32int i, old_size;

	if(size & 0xFFFFF000){
		size &= 0xFFFFF000;
		size += 0x1000;
	}
	if(size < KHEAP_MINSIZE){
		size = KHEAP_MINSIZE;
	}
	old_size = pheap -> end_address - pheap -> start_address;
	for(i = old_size - 0x1000; i >= size; i -= 0x1000)
		free_frame(get_page(pheap -> start_address + i, 0, kernel_directory));
	pheap -> end_address = pheap -> start_address + size;

	return size;
}

static s32int find_smallest_hole(u32int size, u8int align, struct heap *heap){

	u32int iter = 0, location;
	struct header *pheader;
	s32int offset, holesize;
/*
	monitor_write("Request size = ");
	monitor_write_hex(size);
	monitor_write(", ");

	monitor_write("index size = ");
	monitor_write_hex((heap -> index).size);
	monitor_write("\n");
*/
	while(iter < (heap -> index).size){

		pheader = (struct header *)((heap -> index).array[iter]);
/*
		monitor_write("header size = ");
		monitor_write_hex(((struct header *)(heap -> index).array[iter]) -> size);
		monitor_write("\n");
*/
		if(align){
			location = (u32int)pheader;
			offset = 0;
			if((location + sizeof(struct header)) & 0xFFFFF000){
				offset = 0x1000 - (location + sizeof(struct header)) % 0x1000;
			}
			holesize = (s32int)(pheader -> size) - offset;
			if(holesize >= (s32int)size){
				break;
			}
		}else{
			if(pheader -> size >= size){
				break;
			}
		}
		iter++;
	}

	if(iter == (heap -> index).size){
		return -1;
	}else{
		return iter;
	}
}

static u8int header_lt(void *a, void *b){

	return (((struct header *)a) -> size < ((struct header *)b) -> size);
}

struct heap *create_heap(u32int start, u32int end, u32int max, u8int supervisor, u8int readonly){

	struct heap *ph;
	struct header *hole;	

	ph = (struct heap *)kmalloc(sizeof(struct heap));

	ph -> index = place_ordered_map(start, KHEAP_INDEXSIZE, header_lt);

	monitor_write("Placing done.\nPlacement = ");
	monitor_write_hex(placement_addr);
	monitor_write("\n");

	start += sizeof(type_common) * KHEAP_INDEXSIZE;

	monitor_write("heap start addr = ");
	monitor_write_hex(start);
	monitor_write("\n");

	if(start & 0xFFFFF000){
		start &= 0xFFFFF000;
		start += 0x1000;
	}	
	ph -> start_address = start;
	ph -> end_address = end;
	ph -> max_address = max;
	ph -> supervisor = supervisor;
	ph -> readonly = readonly;

	hole = start;
	hole -> magic = KHEAP_MAGIC;
	hole -> is_hole = 1;
	hole -> size = end - start;
	insert_ordered_map(&(ph -> index), (void *)hole);
	
	monitor_write("index size = ");
	monitor_write_hex((ph -> index).size);
	monitor_write("\n");

	return ph;
}

/* alloc() and free() are from JamesM's code */

void *alloc(u32int size, u8int page_align, struct heap *heap){

    u32int new_size = size + sizeof(struct header) + sizeof(struct footer);
    s32int iterator = find_smallest_hole(new_size, page_align, heap);

    if (iterator == -1){
        u32int old_length = heap -> end_address - heap->start_address;
        u32int old_end_address = heap -> end_address;

        expand(old_length + new_size, heap);
        u32int new_length = heap -> end_address-heap -> start_address;

        iterator = 0;
        u32int idx = -1; 
	u32int value = 0x0;
        while (iterator < heap->index.size){
            u32int tmp = (u32int)(heap -> index).array[iterator];
            if (tmp > value){
                value = tmp;
                idx = iterator;
            }
            iterator++;
        }

        if (idx == -1){
            struct header *header = (struct header *)old_end_address;
            header -> magic = KHEAP_MAGIC;
            header -> size = new_length - old_length;
            header -> is_hole = 1;
            struct footer *footer = (struct footer *) (old_end_address + header -> size - sizeof(struct footer));
            footer -> magic = KHEAP_MAGIC;
            footer -> header = header;
            insert_ordered_map(&heap -> index, (void *)header);
        }
        else{
            struct header *header = (heap -> index).array[iterator];
            header -> size += new_length - old_length;
            struct footer *footer = (struct footer *) ((u32int)header + header -> size - sizeof(struct footer));
            footer -> header = header;
            footer -> magic = KHEAP_MAGIC;
        }
        return alloc(size, page_align, heap);
    }

    struct header *orig_hole_header = (struct header *)(heap -> index).array[iterator];
    u32int orig_hole_pos = (u32int)orig_hole_header;
    u32int orig_hole_size = orig_hole_header -> size;
    if (orig_hole_size-new_size < sizeof(struct header) + sizeof(struct footer)){
        size += orig_hole_size - new_size;
        new_size = orig_hole_size;
    }

    if (page_align && orig_hole_pos & 0xFFFFF000){
        u32int new_location   	   	= orig_hole_pos + 0x1000 - (orig_hole_pos & 0xFFF) - sizeof(struct header);
        struct header *hole_header 	= (struct header *)orig_hole_pos;
        hole_header -> size     	= 0x1000 - (orig_hole_pos & 0xFFF) - sizeof(struct header);
        hole_header -> magic    	= KHEAP_MAGIC;
        hole_header -> is_hole  	= 1;
        struct footer *hole_footer 	= (struct footer *) ( (u32int)new_location - sizeof(struct footer) );
        hole_footer -> magic    	= KHEAP_MAGIC;
        hole_footer -> header   	= hole_header;
        orig_hole_pos         		= new_location;
        orig_hole_size        		= orig_hole_size - hole_header->size;
    }
    else{
        delete_ordered_map(&heap->index, iterator);
    }

    struct header *block_header  = (struct header *)orig_hole_pos;
    block_header -> magic     	 = KHEAP_MAGIC;
    block_header -> is_hole   	 = 0;
    block_header -> size      	 = new_size;
    struct footer *block_footer  = (struct footer *) (orig_hole_pos + sizeof(struct header) + size);
    block_footer -> magic     	 = KHEAP_MAGIC;
    block_footer -> header    	 = block_header;

    if (orig_hole_size - new_size > 0){
        struct header *hole_header 	= (struct header *) (orig_hole_pos + sizeof(struct header) + size + sizeof(struct footer));
        hole_header -> magic 		= KHEAP_MAGIC;
        hole_header -> is_hole  	= 1;
        hole_header -> size		= orig_hole_size - new_size;
        struct footer *hole_footer 	= (struct footer *) ((u32int)hole_header + orig_hole_size - new_size - sizeof(struct footer));
        if ((u32int)hole_footer < heap -> end_address){
            hole_footer -> magic = KHEAP_MAGIC;
            hole_footer -> header = hole_header;
        }
        insert_ordered_map(&heap -> index, (void*)hole_header);
    }
    
    return (void *) ((u32int)block_header+sizeof(struct header));
}


void free(void *p, struct heap *heap){

    if (p == 0)
        return;

    struct header *header = (struct header *) ( (u32int)p - sizeof(struct header) );
    struct footer *footer = (struct footer *) ( (u32int)header + header -> size - sizeof(struct footer) );

    header -> is_hole = 1;

    char do_add = 1;

    struct footer *test_footer = (struct footer *) ( (u32int)header - sizeof(struct footer) );
    if (test_footer -> magic == KHEAP_MAGIC && test_footer -> header -> is_hole == 1){
        u32int cache_size = header -> size;
        header = test_footer -> header;
        footer -> header = header;
        header -> size += cache_size;
        do_add = 0;
    }

    struct header *test_header = (struct header *) ( (u32int)footer + sizeof(struct footer) );
    if (test_header -> magic == KHEAP_MAGIC && test_header -> is_hole){
        header -> size += test_header -> size;
        test_footer = (struct footer *) ( (u32int)test_header + test_header -> size - sizeof(struct footer) );
        footer = test_footer;
        u32int iterator = 0;
        while ((iterator < heap -> index.size) && ((heap -> index).array[iterator] != (void*)test_header) )
            iterator++;

        delete_ordered_map(&heap -> index, iterator);
    }

    if ( (u32int)footer + sizeof(struct footer) == heap -> end_address){
        u32int old_length = heap -> end_address - heap -> start_address;
        u32int new_length = contract((u32int)header - heap -> start_address, heap);
        if (header -> size - (old_length - new_length) > 0){
            header -> size -= old_length - new_length;
            footer = (struct footer *) ( (u32int)header + header -> size - sizeof(struct footer) );
            footer -> magic = KHEAP_MAGIC;
            footer -> header = header;
        }else{
            u32int iterator = 0;
            while ( (iterator < heap->index.size) &&  ((heap -> index).array[iterator] != (void*)test_header) )
                iterator++;
            if (iterator < heap->index.size)
                delete_ordered_map(&heap -> index, iterator);
        }
    }

    if (do_add == 1)
        insert_ordered_map(&heap -> index, (void *)header);

}
